import os
from business.utils.vector_store import VectorStore
import sys

def check_vector_store(vector_store_path):
    """检查向量库文件是否存在"""
    faiss_path = os.path.join(vector_store_path, "index.faiss")
    pkl_path = os.path.join(vector_store_path, "index.pkl")
    
    if not os.path.exists(vector_store_path):
        return False, f"向量库目录不存在: {vector_store_path}"
    if not os.path.exists(faiss_path):
        return False, f"向量库索引文件不存在: {faiss_path}"
    if not os.path.exists(pkl_path):
        return False, f"向量库配置文件不存在: {pkl_path}"
    return True, "向量库文件检查通过"

def main():
    # 获取当前工作目录
    current_dir = os.getcwd()
    print(f"当前工作目录: {current_dir}")
    
    # 设置向量存储路径 - 确保和 rag_init.py 使用相同的路径
    vector_store_path = os.path.join(current_dir, "data/vector")
    
    print("向量检索测试程序")
    print("-" * 50)
    print(f"向量库路径: {vector_store_path}")
    
    # 检查向量库文件
    is_valid, message = check_vector_store(vector_store_path)
    if not is_valid:
        print(f"\n错误: {message}")
        print("\n请确保在正确的目录下运行，且已执行 rag_init.py 初始化向量库")
        return
    
    try:
        # 初始化向量存储
        print("\n正在加载向量库...")
        vector_store = VectorStore()
        vs = vector_store.load_vector_store(vector_store_path)
        
        if not vs:
            print("错误：无法加载向量库，请确保已经创建并初始化了向量库")
            return
            
        print("向量库加载成功！")
        
        # 交互式查询循环
        while True:
            print("\n" + "-" * 50)
            query = input("\n请输入您的问题 (输入 'q' 退出): ").strip()
            
            if query.lower() == 'q':
                print("感谢使用！再见！")
                break
                
            if not query:
                print("问题不能为空，请重新输入")
                continue
                
            print("\n正在检索相关内容...")
            results = vs.similarity_search(query, k=3)
            
            print("\n检索结果:")
            for i, doc in enumerate(results, 1):
                print(f"\n=== 结果 {i} ===")
                print(doc.page_content)
                print("-" * 30)

    except Exception as e:
        print(f"发生错误: {e}")
        print(f"错误类型: {type(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main() 